package com.ycy.lucene.plugin;


import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.codecs.DocValuesFormat;
import org.apache.lucene.codecs.lucene80.Lucene80Codec;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.search.spell.LevenshteinDistance;
import org.apache.lucene.search.spell.PlainTextDictionary;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.RAMDirectory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.PostConstruct;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.regex.Pattern;

/**
 * 使用query为语料库进行ngram纠错
 */

public class NgramSpellChecker {
  private static final Logger logger = LoggerFactory.getLogger(NgramSpellChecker.class);

  private static final float ACCURACY = 0.8f;
  private static final int STROKES_MAX_LENGTH = 150;
  private static final int STROKES_MIN_LENGTH = 4;

  public Pattern pattern = Pattern.compile("\\w");
  public CustomSpellChecker wordChecker = null;
  public CustomSpellChecker enameChecker = null;

  PinyinHelper pinyinHelper = new PinyinHelper();

  //存放字母以及笔画对应的原汉字用作纠错结果映射
  public Map<String, String> dicMappingMap = null;

  private Set<String> dicMappingSet = null;

  @PostConstruct
  private void inti() {
    long beginTime=System.currentTimeMillis();
//    loadDicDataToIndex();
    loadEnameDataToIndex();
    getSpellCheckWord("dongfamgmingzhu");
    logger.info("NgramSpellChecker.init cost :{} ",(System.currentTimeMillis()-beginTime));
  }

  public static void main(String[] args) {
    System.out.println(".....");
    long start = System.currentTimeMillis();
    NgramSpellChecker sp = new NgramSpellChecker();
    sp.loadEnameDataToIndex();
    System.out.println(sp.getSpellCheckWord("poncowarno"));
    System.out.println("cost:" + (System.currentTimeMillis() - start));
  }

  public List<String> getENameSpellCheckWord(String keyword) throws IOException {
    if(enameChecker.exist(keyword)) {
      return Collections.emptyList();
    }
    String[] words = keyword.split("\\s");
    StringBuilder correctWord = new StringBuilder();

    boolean checked = true;
    for (String word : words) {
      if (StringUtils.length(word) > 10) {
        checked = false;
        break;
      }
      if(StringUtils.isEmpty(word)) {
        continue;
      }
      String[] suggestions = wordChecker.suggestSimilar(word,4,0.44f);
      if (dicMappingSet.contains(keyword) || wordChecker.exist(word) || suggestions.length == 0) {
        correctWord.append(word);
      } else {
        correctWord.append(suggestions[0]);
      }
      correctWord.append(" ");
    }

    List<String> spellcheckList = Lists.newArrayList();

    if (checked) {
      correctWord.setLength(correctWord.length() - 1);
      if (!StringUtils.equals(correctWord.toString(), keyword)) {
        spellcheckList.add(correctWord.toString());
      }
    }

    if (keyword.length() > 20) {
      return spellcheckList;
    }

    String[] suggestions = wordChecker.suggestSimilar(keyword,2,0.44f);

    spellcheckList.addAll(Arrays.asList(suggestions));
    return spellcheckList;
  }

  private boolean containsWitSpace(String s1, String s2) {
    if (StringUtils.isEmpty(s1) || StringUtils.isEmpty(s2)) {
      return false;
    }
    return StringUtils.contains(s1.replace("\\s", ""), s2.replace("\\s", ""));
  }

  public List<String> getSpellCheckWord(String keyword) {
    List<String> spellCheckList = Lists.newArrayList();
    try {

      List<String> eNameSpellCheckWord = getENameSpellCheckWord(keyword);
      spellCheckList.addAll(eNameSpellCheckWord);


    } catch (Exception e) {
      logger.info("getSpellCheckWord error {} {}", keyword, e);
    }

    return Lists.newArrayList(Sets.newHashSet(spellCheckList));
  }


  public void loadEnameDataToIndex() {
    CustomSpellChecker tmpSpellChecker = null;

    try {
      IndexWriterConfig config = new IndexWriterConfig(new StandardAnalyzer());
      config.setCodec(new Lucene80Codec(){
        @Override
        public DocValuesFormat getDocValuesFormatForField(String field) {
          return DocValuesFormat.forName("Direct");//Direct or Memory
        }
      });
      Directory directory = new RAMDirectory();
      // 采用NGram方式分词纠错
      tmpSpellChecker = new CustomSpellChecker(directory, new LevenshteinDistance());
      tmpSpellChecker.setAccuracy(ACCURACY);
      BufferedReader reader = new BufferedReader(new FileReader(this.getClass().getClassLoader().getResource("data/ngramspellcheck-dic/ename.txt").getPath()));
      tmpSpellChecker.indexDictionary(new PlainTextDictionary(reader), config, true);
    } catch (Exception e) {
      logger.info("load ename spellcheck error fail", e);
    }

    enameChecker = tmpSpellChecker;
  }

}
